@file:UseSerializers(ResourceLocationSerializer::class, ItemStackSerializer::class, TextAsJsonStringSerializer::class)

package com.lalameow.researchstation.researching

import com.lalameow.researchstation.ResearchStation
import com.lalameow.researchstation.api.technology.SingleTechnologyCriterion
import com.lalameow.researchstation.api.technology.Technology
import com.lalameow.researchstation.api.technology.TechnologyRequirement
import com.lalameow.researchstation.api.technology.TechnologyReward
import com.lalameow.researchstation.translatable
import com.lalameow.researchstation.util.serialization.IndirectSerialDescriptor
import com.lalameow.researchstation.util.serialization.TextAsJsonStringSerializer
import kotlinx.serialization.Contextual
import kotlinx.serialization.ExperimentalSerializationApi
import kotlinx.serialization.KSerializer
import kotlinx.serialization.Serializable
import kotlinx.serialization.UseSerializers
import kotlinx.serialization.builtins.MapSerializer
import kotlinx.serialization.encoding.Decoder
import kotlinx.serialization.encoding.Encoder
import net.minecraft.item.ItemStack
import net.minecraft.server.network.ServerPlayerEntity
import net.minecraft.text.Text
import net.minecraft.util.Identifier
import net.silkmc.silk.core.serialization.serializers.ResourceLocationSerializer
import net.silkmc.silk.nbt.serialization.serializer.ItemStackSerializer

fun ResearchingTree(technologies: Set<Technology>, context: ServerPlayerEntity): ResearchingTree {
    val flatTree = ResearchingTree()
    val researching = technologies.associate {
        it.id to it.toResearching(context)
    }
    val criteria =
        technologies.associate { it.id to it.criteria }
    val criterionTechnologies =
        technologies.associate { technology ->
            technology.id to technology.criteria
                .filterIsInstance<SingleTechnologyCriterion>()
                .mapTo(mutableSetOf()) { it.technology.id }
        }
    val invisible = technologies.filterTo(mutableSetOf()) { technology ->
        // 不可见且不满足任一前置
        !technology.visible && technology.criteria.any { it.test(context) != null }
    }.mapTo(mutableSetOf()) { it.id }

    fun ResearchingTechnology.tree(): ResearchingTree? {
        // 当前科技是否可见
        var visible = true
        val currentTrees = flatTree[this] ?: ResearchingTree()
        // 遍历当前科技的所有前置科技
        for (criterionTechnology in criterionTechnologies[id]!!) {
            if (criterionTechnology in invisible) {
                visible = false
                break
            }
            // 前置科技对应树
            val criterionTree = researching[criterionTechnology]!!.tree()
            // 前置科技是否可见
            visible = criterionTree != null
            if (!visible) break
            // 添加当前科技到前置节点中
            criterionTree?.set(this, currentTrees)
        }
        if (visible) {
            flatTree[this] = currentTrees
            return currentTrees
        }
        return null
    }

    for (technology in researching.filterKeys { it !in invisible }) technology.value.tree()
    return ResearchingTree(
        flatTree
            .filterKeys { criteria[it.id].isNullOrEmpty() } as LinkedHashMap<ResearchingTechnology, ResearchingTree>
    )
}

@Serializable
data class ResearchingTechnology(
    val id: Identifier,
    val icon: Set<ItemStack>,
    /**
     * [Can't use wildcard](https://github.com/Kotlin/kotlinx.serialization/issues/607)
     */
    val requirements: Set<TechnologyRequirement<out @Contextual Any?>>,
    val rewards: Set<TechnologyReward>,
    val criteria: Set<Text>,
    val name: Text = ResearchStation.translatable("technology.$id"),
    val description: Text? = null
) {
    constructor(technology: Technology, context: ServerPlayerEntity) : this(
        technology.id,
        technology.icon,
        technology.requirements,
        technology.rewards,
        technology.criteria.mapNotNullTo(mutableSetOf()) { it.test(context) },
        technology.name,
        technology.description
    )
}

private fun Technology.toResearching(context: ServerPlayerEntity) = ResearchingTechnology(this, context)

@Serializable(ResearchingTree.Serializer::class)
class ResearchingTree(
    map: LinkedHashMap<ResearchingTechnology, ResearchingTree> = LinkedHashMap()
) : LinkedHashMap<ResearchingTechnology, ResearchingTree>(), MutableMap<ResearchingTechnology, ResearchingTree> by map {
    @OptIn(ExperimentalSerializationApi::class)
    object Serializer : KSerializer<ResearchingTree> {
        private val serializer by lazy { MapSerializer(ResearchingTechnology.serializer(), this) }

        override val descriptor = IndirectSerialDescriptor { serializer.descriptor }

        override fun deserialize(decoder: Decoder) =
            ResearchingTree(serializer.deserialize(decoder) as LinkedHashMap) // MapSerializer 总是 LinkedHashMapSerializer

        override fun serialize(encoder: Encoder, value: ResearchingTree) = serializer.serialize(encoder, value)
    }
}
